#include "common.h"
#include "memory/cache.h"
#include "memory/memory.h"
#include <stdlib.h>
#include <time.h>

/* Declare DRAM read/write functions */
uint32_t dram_read(hwaddr_t, size_t);
void dram_write(hwaddr_t, size_t, uint32_t);

/* L1 Cache */
static Cache l1_cache;

/* Extract address fields */
static inline uint32_t get_block_offset(hwaddr_t addr) {
	return addr & BLOCK_OFFSET_MASK;
}

static inline uint32_t get_set_index(hwaddr_t addr) {
	return (addr >> BLOCK_OFFSET_BITS) & SET_INDEX_MASK;
}

static inline uint32_t get_tag(hwaddr_t addr) {
	return (addr >> (BLOCK_OFFSET_BITS + SET_INDEX_BITS)) & TAG_MASK;
}

static inline hwaddr_t get_block_addr(hwaddr_t addr) {
	return addr & ~BLOCK_OFFSET_MASK;
}

/* Initialize cache */
void init_cache() {
	int i, j;
	for (i = 0; i < NR_CACHE_SETS; i++) {
		for (j = 0; j < CACHE_WAY; j++) {
			l1_cache.sets[i].lines[j].valid = false;
			l1_cache.sets[i].lines[j].tag = 0;
		}
	}
	/* Initialize random seed */
	srand(time(NULL));
}

/* Find a cache line in the set */
static CacheLine* find_cache_line(uint32_t set_idx, uint32_t tag) {
	int i;
	CacheSet *set = &l1_cache.sets[set_idx];
	
	for (i = 0; i < CACHE_WAY; i++) {
		if (set->lines[i].valid && set->lines[i].tag == tag) {
			return &set->lines[i];
		}
	}
	return NULL;
}

/* Find an invalid cache line or select one for replacement */
static CacheLine* allocate_cache_line(uint32_t set_idx, uint32_t tag) {
	int i;
	CacheSet *set = &l1_cache.sets[set_idx];
	
	/* First, try to find an invalid line */
	for (i = 0; i < CACHE_WAY; i++) {
		if (!set->lines[i].valid) {
			set->lines[i].valid = true;
			set->lines[i].tag = tag;
			return &set->lines[i];
		}
	}
	
	/* All lines are valid, randomly select one for replacement */
	i = rand() % CACHE_WAY;
	set->lines[i].tag = tag;
	return &set->lines[i];
}

/* Cache read operation */
uint32_t cache_read(hwaddr_t addr, size_t len) {
	uint32_t offset = get_block_offset(addr);
	uint32_t set_idx = get_set_index(addr);
	uint32_t tag = get_tag(addr);
	hwaddr_t block_addr = get_block_addr(addr);
	
	/* Check if data crosses cache block boundary */
	if (offset + len > CACHE_BLOCK_SIZE) {
		/* Split into two cache accesses */
		uint32_t len1 = CACHE_BLOCK_SIZE - offset;
		uint32_t len2 = len - len1;
		uint32_t data1 = cache_read(addr, len1);
		uint32_t data2 = cache_read(addr + len1, len2);
		return data1 | (data2 << (len1 * 8));
	}
	
	/* Find cache line */
	CacheLine *line = find_cache_line(set_idx, tag);
	
	if (line != NULL) {
		/* Cache hit */
		uint32_t ret = 0;
		memcpy(&ret, line->data + offset, len);
		return ret;
	} else {
		/* Cache miss - allocate a line and load from DRAM */
		line = allocate_cache_line(set_idx, tag);
		
		/* Load entire block from DRAM */
		int i;
		for (i = 0; i < CACHE_BLOCK_SIZE; i += 4) {
			uint32_t word = dram_read(block_addr + i, 4);
			memcpy(line->data + i, &word, 4);
		}
		
		/* Return requested data */
		uint32_t ret = 0;
		memcpy(&ret, line->data + offset, len);
		return ret;
	}
}

/* Cache write operation */
void cache_write(hwaddr_t addr, size_t len, uint32_t data) {
	uint32_t offset = get_block_offset(addr);
	uint32_t set_idx = get_set_index(addr);
	uint32_t tag = get_tag(addr);
	
	/* Check if data crosses cache block boundary */
	if (offset + len > CACHE_BLOCK_SIZE) {
		/* Split into two cache accesses */
		uint32_t len1 = CACHE_BLOCK_SIZE - offset;
		uint32_t len2 = len - len1;
		uint32_t data1 = data & ((1u << (len1 * 8)) - 1);
		uint32_t data2 = data >> (len1 * 8);
		cache_write(addr, len1, data1);
		cache_write(addr + len1, len2, data2);
		return;
	}
	
	/* Write through: always write to DRAM */
	dram_write(addr, len, data);
	
	/* Find cache line (not write allocate) */
	CacheLine *line = find_cache_line(set_idx, tag);
	
	if (line != NULL) {
		/* Cache hit - update cache line */
		memcpy(line->data + offset, &data, len);
	}
	/* Cache miss - do nothing (not write allocate) */
}

